[OMNIML-6054] Route load_pretrained by the distribution plan - #2779

What does this PR do?

Type of change: New feature

load_pretrained (#2733) loads a model whole. A distributed load, such as parallel_load_and_prepare_fsdp2 (main), is a separate entry point, and every caller picks one itself, as hf_ptq.py --use_fsdp2 does. A caller that picks a loader also gets what that loader skips: a sharded load cannot run the model type's loading hook, so a Nemotron-H checkpoint's MTP head is not built and its tensors go to the export unquantized.

This PR makes load_pretrained the one loader, whole or distributed:

• distribution is a DistributionPlan; by default it is the plan resolve_distribution_plan ([OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737) decides.
  • dense="fsdp": parallel_load_and_prepare_fsdp2 over the process group. A single process launched with torchrun --nproc_per_node=1 has no group, so the load sets one up, of one rank, once nothing can refuse the load anymore; a process started without a launcher cannot set one up and is refused. distributed_run ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736) destroys it with the run. One rank of a multi-rank launch whose group is not set up yet is refused too, told to set it up first with distributed_run(), rather than shard over the launch's ranks under a one-rank plan; the plan-size refusal says the same when that is why the sizes differ.
  • Otherwise from_pretrained. On a distributed run every rank gets the whole model, on its own GPU.
  • A plan that shards the routed experts raises here. [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 adds that load, once the expert-parallel loader ([OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 and [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743) is in. Until it lands, the changelog entry says so: an MoE planned expert-parallel, as the default MODELOPT_EP_DEGREE=auto plans one whose config has a base_model_ep_plan, is refused before it loads; it loads with MODELOPT_EP_DEGREE=1 (with MODELOPT_FSDP_SIZE=world to FSDP2-shard it) or a plan with ep_size=1.
  • resolve_distribution_plan takes the class the load builds with (auto_class, which load_pretrained passes). With a class the sharded loads do not build and no knob asking to shard, it plans the whole model on every rank without a meta-device build, so an encoder-decoder loaded through AutoModel, for example, loads multi-rank as before.
• The plan is recorded on the model and every submodule (set_distribution_plan, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), where calibration and export read it. load_pretrained records it exactly once, at the end of every load; the loaders it calls record none. fsdp2_wrap records nothing, as on main, so a caller that wraps or loads a model through it directly records the plan itself, as the error for a missing plan says. hf_ptq.py --use_fsdp2 loads through parallel_load_and_prepare_fsdp2 directly and records no plan in this PR; nothing on that path reads the plan before [OMNIML-6071] Calibrate a distributed model on per-rank data #2740, which records it in the same PR that makes get_max_batch_size read it.
• fsdp2_wrap moves to model_load_utils, with the load that calls it and every other loader, and imports LayerActivationCollector at module level. modelopt.torch.utils.distributed.fsdp2_wrap forwards to it, so callers of the old path keep working, given the hf extra the loaders need, and the FSDP2 helpers it uses stay in distributed.
• Every load makes a Hub ID local and records the unplaced source keys ([OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733). The FSDP2 load builds its model from the config without the Hub options, so custom code an auto_map names comes with the environment's Hub settings.
• Every rank loads the same plan. Every rank takes rank 0's distribution, a plan or none, through one broadcast (broadcast_from, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736) before it branches, and loads a plan rank 0 was handed just as it would one it decided. Ranks that each loaded the plan they were handed could build mismatched meshes and hang, and ranks passing different kinds would wait in each other's collectives. Rank 0 checks what it was handed before it sends it, a non-plan included, so a distribution it refuses raises on every rank instead of leaving the others waiting in the broadcast.
• Refusals, before anything loads.
  • A sharded load cannot run the model type's loading hook, so it refuses a checkpoint the hook would build modules for: the MTP head above.
  • It refuses from_pretrained's placement arguments (device_map, max_memory, ...) and a model class it does not build, rather than ignore them, and an explicit dtype, saying why it does not apply: the load keeps the checkpoint's dtype, and FSDP2's mp_policy sets the dtype it computes in.
  • The FSDP2 load refuses a quantized checkpoint, whose tensors it would read as stored.
  • The FSDP2 load refuses a VLM, which it would build through AutoModelForCausalLM: refused deep in the loader, or its language model alone. It refuses any config _auto_class does not build through AutoModelForCausalLM, before it sets up a group or reads anything, and the refusal says how to load the model whole on every rank: its model class as auto_class, with MODELOPT_FSDP_SIZE=1 or a replicated plan.
  • A plan made for another world size is refused, and so is anything other than a plan.
  • hf_config without the plan it was resolved with is refused. A sharded load builds the model from it as given, so a plan decided here, from the checkpoint's own config, would not be the one it came with. config= is for a config the plan should be decided from. resolve_distribution_plan's docstring says to hand the config it returns to load_pretrained only with a plan that shards the model: a whole-model load reads the checkpoint's config itself and refuses hf_config.

The FSDP2 route's refusal names its ways out (MODELOPT_FSDP_SIZE=1, or a replicated plan), and a config handed with a plan is copied before the load builds from it. With a plan handed and no config, the FSDP2 route reads the checkpoint's config once and hands it to the load, rather than reading it twice. Deciding a plan builds the meta-device model from a copy of the config, so the config returned with the plan, which a sharded load builds from, is the same on every rank: the build writes its dtype and its attention and experts implementations, and only rank 0 builds.

Usage

# torchrun --nproc_per_node 8 quantize.py
from modelopt.torch.utils.distributed import distributed_run
from modelopt.torch.utils.plugins.model_load_utils import load_pretrained

with distributed_run():
    model = load_pretrained("meta-llama/Llama-3.1-8B")  # a whole copy on every rank's GPU
    # With MODELOPT_FSDP_SIZE=world: FSDP2-sharded over the 8 ranks instead.

Testing

• New tests/unit/torch/utils/test_load_pretrained_distribution.py. It runs on CPU, and stand-in loaders record their calls. It covers:
  • each load: the arguments it gets and the plan it records;
  • a decided FSDP2 load, handed the config the plan was decided from instead of reading it again, and a handed FSDP2 plan's load reading the checkpoint's config once (that check fails without its fix);
  • a single process's FSDP2 load setting up its group of one, and leaving none behind when the load is refused;
  • every refusal, raised before any loader runs, the expert-parallel plan's included; test_a_load_refuses_what_it_cannot_load has a VLM row (a Llava config), which fails without the refusal;
  • on two real gloo ranks (cpu_dist_workers), rank 0 handed a replicated plan and rank 1 an FSDP2 one, or none: both load whole and record rank 0's plan. A rank that loads the plan it was handed fails at once, on the placement refusal. Rank 0 handed a non-plan fails both ranks;
  • a caller's own model class: loaded whole on every rank, planned without building the model, with no sharding knob, and refused by the sharded load with one.
• tests/unit/torch/utils/test_resolve_distribution_plan.py: the config returned with the plan is the one every rank read, unchanged by rank 0's meta-device build (new test_the_config_returned_with_the_plan_is_as_every_rank_read_it, which fails without its fix).
• A single process's load, placed as from_pretrained places it, with the caller's arguments and a single process's plan: [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's real CPU load in tests/unit/torch/utils/test_unplaced_source_keys.py, which now also checks the plan.
• fsdp2_wrap leaving the plan to the load: a CPU unit test in tests/unit/torch/utils/test_load_pretrained_distribution.py. [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's GPU tests that wrap or load through it directly (test_export_carry_over.py, test_model_load_utils.py) record the plan themselves. Its old path forwarding every argument: tests/unit/torch/utils/test_distributed_run.py.
• Mutation-checked: 17 mutants, each failing a test. They drop the rank-GPU default, pin a single process's model, drop the config handoff or the world-size check, load an expert-parallel or FSDP2 plan whole, accept a string, let each rank load the plan it was handed, record no plan in the load or record one in fsdp2_wrap too, set no group up for a single process's FSDP2 load, set one up over a live group, or have the old fsdp2_wrap path drop an argument or the result.
• New tests/gpu/torch/utils/test_load_pretrained_distribution.py runs load_pretrained left to decide, on every GPU of the node:
  • a dense model loads whole on each rank, and with MODELOPT_FSDP_SIZE=world FSDP2-sharded;
  • each rank's weights match the checkpoint;
  • the load records its unplaced keys and its plan.
  • Not run locally; CI's GPU workflow runs it.
• tests/examples/hf_ptq/test_example_utils.py: get_model's stand-in model is a module now, since load_pretrained records its plan on it.
• The load refusals are one parametrized test, and the single-process load test gives way to [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's real load. Each drop is checked by mutation.
• A single process without a launcher is refused an FSDP2 load, without stubbing the setup, and a handed config is left as the caller passed it. A rank of a multi-rank launch whose process group is not set up is refused any load, before anything is fetched. The refusal table has an explicit dtype.
• Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.
• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.
• Review fixes, 2026-10-27: a rank of a multi-rank launch without its process group is refused on every route, before anything is fetched; a caller's model class with a knob asking to shard is refused before the model is built; a changelog entry for fsdp2_wrap's move. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.
• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739's FSDPModule import, now on main, and this PR's CPUOffloadPolicy, fully_shard are one import line, and fsdp2_wrap sits ahead of the meta-DCP loader, as before.
• Review fixes, 2026-10-27 (Claude's review after the rebase): a sharded load takes dtype=None and torch_dtype as from_pretrained does. None and "auto" keep the checkpoint's dtype under either name, and any other dtype is refused with the dtype message, where torch_dtype got the placement one. The new cases fail without the change.
• Review suggestions, 2026-10-27 evening: a launch whose processes set no group up (hf_ptq under torchrun without --use_fsdp2) loads the whole model on each process again, refusing there only a plan for several ranks and a sharded load; a handed plan is checked before the checkpoint is fetched; an FSDP2 layout decided from rank 0's GPU memory says so when it refuses placement options; fsdp2_wrap imports the layer detection where it uses it.
• Review fixes, 2026-10-27 night: a launch without its process group refuses MODELOPT_EP_DEGREE above 1 or MODELOPT_FSDP_SIZE naming the world before anything is fetched, and the sharded-load refusal names the plan to pass. The refusal test's new rows fail without the change.
• Review suggestions, 2026-10-27, last round: on a launch without its group, a MODELOPT_EP_DEGREE that is not an EP degree is refused with the plan decision's message, and a padded value is read as its number; load_pretrained's docstring says the automatic layout weighs the model in the checkpoint's own dtype. The refusal test's new row fails without the change.
• Review fix, 2026-10-27, last round: a one-rank sharded plan in a launch without its group is refused before the checkpoint is fetched, and the advice for a model with no plan names load_pretrained again, which records the plan from this PR on.
• Review fixes, 2026-10-27, last round: hf_config with a plan that loads the model whole is refused before the checkpoint is fetched, with the no-plan case's message, and the new pre-fetch test row fails without the change; fsdp2_wrap's docstring says the FSDP2 load wraps with it.
• Review suggestions, 2026-10-28: the refusals that depend only on a handed plan's arguments (auto_class, dtype, options the sharded load does not take) run before the checkpoint is fetched, and the new test row fails without the change; a test docstring no longer says fsdp2_wrap loads.
• Review suggestions, 2026-10-28, fifth round: the handed-plan test pins the auto_class and dtype arms of the pre-fetch refusal, each failing with its arm left out of the pre-fetch call; the four function-local imports of modelopt.torch.utils.distributed go, now that [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 imports it at module top.
• Declined review suggestion, 2026-10-28, fifth round: one helper for the kwargs view the pre-fetch refusal and _load_sharded read; the two agree, and the handed-plan test pins each pre-fetch arm.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ load_pretrained keeps its behavior in a single process unless the plan shards the model (FSDP2 there needs torchrun --nproc_per_node=1), and distribution is new. fsdp2_wrap behaves as on main and records no plan; its old import path still works, given the hf extra.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ 0.49.0, Misc.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 6 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2737. The expert-parallel half of this routing is #2780.

Replaces #2774, which carried this routing whole as [19/20]. Moving its branch down the stack made the branch part of #2774's old base, so GitHub closed #2774 as merged into that base; nothing reached main.

Summary by CodeRabbit

• New Features

  • Added GGML imatrix calibration with weighted search and per-expert statistics for fused MoE models.
  • Updated IQ recipes for combined GPTQ and imatrix calibration, and added a Nemotron 3.5 Super VL AutoQuantize recipe.
  • Added standalone draft model training and export for external-draft speculative decoding.
  • Distributed loading supports replicated or FSDP2-sharded models through configurable distribution plans, which are recorded on loaded models. Automatic plan selection and Hub authentication and caching options are supported.

• Documentation

  • Added guidance for distributed loading and distribution-plan recording.